[RLlib] Add gradient checks to avoid nan gradients in TorchLearner.
#47452
+52
−1
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Signed-off-by: simonsays1980 <[email protected]>
…in highly unstable training phases. This helps to keep the optimizer's internal state intact whoch could get corrupted with many zero gradients. Furthermore, added better logging messages. Signed-off-by: simonsays1980 <[email protected]>
simonsays1980
requested review from
sven1977 and
ArturNiederfahrenhorst
as code owners
Signed-off-by: simonsays1980 <[email protected]>
simonsays1980
changed the title
nan gradients in TorchLearner.
sven1977
reviewed
Sep 2, 2024
| @@ -176,7 +176,27 @@ def compute_gradients( | |||
| def apply_gradients(self, gradients_dict: ParamDict) -> None: | |||
| # Set the gradient of the parameters. | |||
| for pid, grad in gradients_dict.items(): | |||
| self._params[pid].grad = grad | |||
| # If updates should not be skipped turn `nan` gradients to zero. | |||
There was a problem hiding this comment.
Wait, I'm confused. We have this block here further below, which I think does the exact same thing: skips the entire optim.step in case any gradient is non-finite (inf or nan).
# `step` the optimizer (default), but only if all gradients are finite.
elif all(
param.grad is None or torch.isfinite(param.grad).all()
for group in optim.param_groups
for param in group["params"]
):
Can you check and see whether these two logics can be consolidated?
Kind of like this:
- If user sets this flag (default=False), the optimizer will skip the update step entirely (+ warning raised by RLlib).
- If user does NOT set this flag (default behavior), grads that are non-finite will be set to 0.0 (+ warning raised by RLlib).
sven1977
changed the title
nan gradients in TorchLearner.nan gradients in TorchLearner.
…lution considers non-finite gradients and gives the user still the option to set such gradients to zero to keep the optimizer's internal state intact. Signed-off-by: simonsays1980 <[email protected]>
sven1977
approved these changes
Sep 3, 2024
ujjawal-khare
pushed a commit
to ujjawal-khare-27/ray
that referenced
this pull request
Oct 15, 2024
…`. (ray-project#47452) Signed-off-by: ujjawal-khare <[email protected]>
ujjawal-khare
pushed a commit
to ujjawal-khare-27/ray
that referenced
this pull request
Oct 15, 2024
…`. (ray-project#47452) Signed-off-by: ujjawal-khare <[email protected]>
ujjawal-khare
pushed a commit
to ujjawal-khare-27/ray
that referenced
this pull request
Oct 15, 2024
…`. (ray-project#47452) Signed-off-by: ujjawal-khare <[email protected]>
ujjawal-khare
pushed a commit
to ujjawal-khare-27/ray
that referenced
this pull request
Oct 15, 2024
…`. (ray-project#47452) Signed-off-by: ujjawal-khare <[email protected]>
ujjawal-khare
pushed a commit
to ujjawal-khare-27/ray
that referenced
this pull request
Oct 15, 2024
…`. (ray-project#47452) Signed-off-by: ujjawal-khare <[email protected]>
ujjawal-khare
pushed a commit
to ujjawal-khare-27/ray
that referenced
this pull request
Oct 15, 2024
…`. (ray-project#47452) Signed-off-by: ujjawal-khare <[email protected]>
ujjawal-khare
pushed a commit
to ujjawal-khare-27/ray
that referenced
this pull request
Oct 15, 2024
…`. (ray-project#47452) Signed-off-by: ujjawal-khare <[email protected]>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Why are these changes needed?
If any gradients turn
naninTorchLearnerthese gradients get added to the network's weights and in turn weights becomenanand all network outputs as well. As a result the training errors out and stops. This PR proposes a gradient check to only add gradients if they are sane. It switchesnanvalues to zeros in the gradients or skips en update entirely.The latter can be of advantage, if training phase ecnounters highly unstable policy updates (e.g. with highly explorative policies or during early stages of training). In such phases many gradients could turn
nanand this may lead to corrupted internal optimizer states (e.g. Adam).Related issue number
#47451
Checks
git commit -s) in this PR.scripts/format.shto lint the changes in this PR.method in Tune, I've added it in
doc/source/tune/api/under thecorresponding
.rstfile.